
rm(list=ls())

jobnum=1

load("parsedData-servantFlanker.Rdata")

newSeed=Sys.time()
set.seed(as.numeric(newSeed))

data=data[[jobnum]]





rdata=function(n,x,par.names) {
  names(x)=par.names
  out=list(Cond=NULL,Resp=NULL,Time=NULL)
  
  for (cond in conds) {
    for (stim in stims) {
      N=as.double(n[cond,stim])
      
      params=x
      params.names=c(names(params),"stoch.s")
      params=c(params,0.1)
      if (stim=="left") {
        params["z"]=x["z"]
      } else if (stim=="right") {
        params["z"]=(1-x["z"])
      }
      if (cond=="comp") {
        congruency=1
      } else if (cond=="incomp") {
        congruency=-1
      }
      
      params["ster"]=x["ster"]*x["ter"]*2
      
      names(params)=params.names
      
      tmp=simulate.DSTP(N=N,params=params,maxCounter=use.maxCounter,stepSize=stepSize,congruency=congruency,use.table=use.table,n.table.options=n.table.options)
      
      out$Resp=c(out$Resp,tmp$resp)
      out$Time=c(out$Time,tmp$rt)
      out$Cond=c(out$Cond,rep(cond,n[cond,stim]))
      out$Stim=c(out$Stim,rep(stim,n[cond,stim]))
    }
  }
  out
}

log.dens.like=function(x,data,par.names){
  names(x)=par.names
  out=0
  
  for (cond in conds) {
    for (stim in stims) {
      
      tmp=data$Cond==cond & data$Stim==stim
      
      I=IQR(data$Time[tmp])
      S=sd(data$Time[tmp])
      bandwidth=0.9*min(I,S)*(Nsamples)^(-0.2)
      
      params=x
      params.names=c(names(params),"stoch.s")
      params=c(params,0.1)
      if (stim=="left") {
        params["z"]=x["z"]
      } else if (stim=="right") {
        params["z"]=(1-x["z"])
      }
      if (cond=="comp") {
        congruency=1
      } else if (cond=="incomp") {
        congruency=-1
      }
      
      params["ster"]=x["ster"]*x["ter"]*2
      
      names(params)=params.names
      
      tmp1=Log.likelihood.fun(all.data=list(rt=data$Time[tmp],resp=data$Resp[tmp]),params=params,Nsample=Nsamples,maxCounter=use.maxCounter,stepSize=stepSize,congruency=congruency,bandwidth=bandwidth,getSimData=simulate.DSTP)
      
      if (length(tmp1[[1]]) == 1) {
        tmp2=tmp1[[1]][[1]]
      } else {
        tmp2=c(tmp1[[1]][[1]],tmp1[[1]][[2]])
      }
      
      out=out+sum(log(pmax(tmp2,1e-10)))
      
      out=out+pmax(dbinom(sum(data$Resp[tmp]==-1),length(data$Resp[tmp]),tmp1[[2]]/Nsamples,log=TRUE),log(1e-10))
    }
  }
  out
}

log.dens.prior=function(x,hyper){
  out=0
  for (p in names(x)) out =
      out+dtnorm(x[p],hyper[[p]][1],hyper[[p]][2],lower.bounds[p],upper.bounds[p],log=TRUE)
  out
}

crossover=function(i,pars,use.theta,use.like,data,hyper,par.names,currIT){
  if (currIT %% 5 == 0) use.like[i] = log.dens.like(use.theta[i,],data,par.names=par.names)
  use.weight=use.like[i]*temperatures[i] + log.dens.prior(use.theta[i,],hyper)
  gamma = 2.38/sqrt(2*length(pars))
  index=sample(c(1:n.chains)[-i],2,replace=F)
  theta=use.theta[i,]						
  theta[pars]=use.theta[i,pars] + gamma*(use.theta[index[1],pars]-use.theta[index[2],pars]) + runif(1,-b,b)
  prior.like=log.dens.prior(theta,hyper)
  if (prior.like > -Inf) {
    like=log.dens.like(theta,data,par.names=par.names)
  } else {
    like = -Inf
  }
  weight=like*temperatures[i] + prior.like
  if(!is.finite(weight))weight=-Inf
  if(runif(1) < exp(weight-use.weight)) {							
    use.theta[i,]=theta
    use.like[i]=like
  }
  c(use.like[i],use.theta[i,])
}

migration.crossover=function(pars,use.theta,use.like,data,hyper,par.names){
  n.migration.chains=ceiling(runif(1,0,n.chains))
  use.chains=sample(1:n.chains,n.migration.chains)
  migration.use.weight=rep(NA,n.migration.chains)
  migration.weight=rep(NA,n.migration.chains)
  for (mi in 1:n.migration.chains) {
    migration.use.weight[mi]=use.like[use.chains[mi]]*temperatures[i] + log.dens.prior(use.theta[use.chains[mi],pars],hyper)
    newChain = mi - 1
    if (newChain == 0) newChain = n.migration.chains
    migration.weight[mi]=use.like[use.chains[newChain]]*temperatures[i] + log.dens.prior(use.theta[use.chains[newChain],pars],hyper)
    if(runif(1) < exp(migration.weight[mi]-migration.use.weight[mi])) {      				
      use.theta[use.chains[mi],]=use.theta[use.chains[newChain],]
      use.like[use.chains[mi]]=use.like[use.chains[newChain]]
    }
  }
  cbind(use.like,use.theta)
}



library(msm)
source("../Code/simulate-DSTP.R")
source("../Code/PDA.R")
dyn.load("../Code/DSTP.so")

stepSize=0.01

use.maxCounter=150




for (use.run in 1:5) {
  
  
  
K = 40
alpha = 0.3
all.temperatures = rev((0:(K-1)/(K-1))^(1/alpha))
temperatures=all.temperatures





start.points=c(.3,.1,.3,0.5,.1,.1,.4,1,0.1,.2,0.1)
lower.bounds=c(-Inf,0,0,0,0,-Inf,-Inf,-Inf,0,0,0)
upper.bounds=c(Inf,Inf,Inf,1,Inf,Inf,Inf,Inf,1,Inf,1)


n.chains=K
n.pars=11
nmc=5000
Nsamples=10000

migration.start=nmc
migration.end=nmc
migration.freq=nmc



b=.001

theta=array(NA,c(n.chains,n.pars,nmc))
weight=array(-Inf,c(nmc,n.chains))


theta.names = c("vT", "a", "ter", "z", "c", "vA","vSS","vRS2","ster","sv","sz")
colnames(theta) = theta.names

names(lower.bounds) = names(upper.bounds) = theta.names



for(i in 1:n.chains){
  current.like=weight[1,i]
  while (current.like==-Inf) {
    current.thetas=rtnorm(n=n.pars,mean=start.points,sd=start.points/5,lower.bounds,upper.bounds)
    current.weight=log.dens.like(current.thetas,data=data,par.names=theta.names)
    current.like=current.weight*temperatures[i]
  }
  theta[i,,1]=current.thetas
  weight[1,i]=current.weight
}


prior=list()

tmp=grep("c",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.2,.2)
}

tmp=grep("a",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.2,.2)
}

tmp=grep("ter",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.3,.3)
}

tmp=grep("z",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.5,.5)
}

tmp=grep("vT",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.3,.3)
}

tmp=grep("vA",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.2,.2)
}

tmp=grep("vSS",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.4,.4)
}

tmp=grep("vRS2",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(1.2,1.2)
}

tmp=grep("ster",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.1,.1)
}

tmp=grep("sv",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.2,.2)
}

tmp=grep("sz",theta.names,value=TRUE)
for (n in 1:length(tmp)) {
  tmp2=tmp[n]
  prior[[tmp2]]=c(.1,.1)
}


savefile=paste("TIDE_Fits/DSTP-TIDE_IS",nmc,"servantFlanker_complex_sub-",jobnum,"run_",use.run,sep="-")

savefile=paste(savefile,"Rdata",sep=".")

begin = date()
for(i in 2:nmc){
  cat("\n ",i,"  ")
  if (i %% migration.freq == 0 & i > migration.start & i < migration.end) {
    temp=migration.crossover(pars=1:n.pars,use.theta=theta[,,i-1],use.like=weight[i-1,],data=data,hyper=prior,par.names=theta.names)
  } else {
    temp=t(sapply(1:n.chains,crossover,pars=1:n.pars,use.theta=theta[,,i-1],use.like=weight[i-1,],data=data,hyper=prior,par.names=theta.names,currIT=i))
  }
  weight[i,]=temp[,1]
  theta[,,i]=temp[,2:(n.pars+1)]
}
end = date()
begin
end



save.image(savefile)





thermo_int = function(t,lp,var){
  
  I = sum(sapply(1:(length(t)-1),function(x) (t[x+1] - t[x]) * ((lp[x] + lp[x+1])/2)),na.rm=T)
  cor = sum(sapply(1:(length(t)-1),function(x) (((t[x+1] - t[x])^2)/12) * (var[x+1] - var[x])),na.rm=T)
  
  return(list(ti=I,ti_cor=I-cor))
}


mean.log.like=apply(weight[2001:5000,],2,mean)
var.log.like=apply(weight[2001:5000,],2,var)


TI.TIDE=thermo_int(t=rev(all.temperatures),lp=rev(mean.log.like),var=rev(var.log.like))




save.image(savefile)

}

